import matplotlib.pyplot as plt
import numpy as np
from scipy import signal
from  scipy.fftpack import fft
from MyFFT import FFT
def createSource(freq):
    t = np.linspace(0, 1, freq, False)
    sig = 1.5 * np.sin(10 * np.pi * t) + 2 * np.sin(20 * 2 * np.pi * t) - np.cos(
        80 * np.pi * t)  # 构造10hz和20hz个信号
    return [sig.tolist(),t]
source=createSource(100)
FFTsource=FFT(source[0])
# FFTsource.plot()


def filter(source,type,band):
    # TODO:ilter to be finished
    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)
    ax1.plot(source[1], source[0])
    ax1.set_title('before filtering')
    ax1.axis([0, 1, -2, 2])


    sos = signal.butter(10, band, type, fs=100, output='sos',analog=False)  # 采样率为1000hz，带宽为15hz，输出sos
    filtered = signal.sosfilt(sos,source[0])  # 将信号和通过滤波器作用，得到滤波以后的结果。在这里sos有点像冲击响应，这个函数有点像卷积的作用。
    ax2.plot(source[1], filtered)
    ax2.set_title('filtered')
    fft2=FFT(filtered.tolist())
    fft2.plot()
    plt.tight_layout()
    plt.show()

filter(source,type='lp',band=10)
filter(source,type='hp',band=30)
filter(source,type='bp',band=[10,30])
filter(source,type='bs',band=[10,30])
# filter(source,type='bandpass')
# filter(source,type='bandstop')

# wez=[1,2,3,4,5,6,7,8]
# to2L(wez)
# yy= fft(wez)
# sortReverse(wez)
# mfft(wez)

# arrayAbs(wez)
# arrayAbs(yy)
# fig,axs=plt.subplots(2,1)
# axs[0].plot(range(0,8),wez)
# axs[0].grid(True)
# axs[1].plot(range(0,8),yy)
# axs[1].grid(True)
# fig.tight_layout()
# plt.show()


